package net.lab1024.sa.ai.stream.module.business.ai.integration;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import net.lab1024.sa.ai.stream.module.business.ai.entity.AiModelConfigEntity;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.HttpClientErrorException;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import java.util.*;

@Service
public class DeepseekAiService implements AiModelAdapter {
    private static final String MODEL_NAME = "deepseek";
    @Autowired
    private LlmErrorCodeHandlerFactory errorCodeHandlerFactory;

    @Override
    public String chat(String prompt, Map<String, Object> params, AiModelConfigEntity config) {
        try {
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", config.getModelType() != null ? config.getModelType() : "deepseek-chat");
            List<Map<String, String>> messages = new ArrayList<>();
            if (params != null && params.containsKey("messages")) {
                // 兼容已有messages
                messages = (List<Map<String, String>>) params.get("messages");
            } else {
                messages.add(Map.of("role", "user", "content", prompt));
            }
            requestBody.put("messages", messages);
            if (params != null) {
                requestBody.putAll(params);
            }
            // 日志：请求体
            System.out.println("[Deepseek] Request Body: " + JSON.toJSONString(requestBody));

            RestTemplate restTemplate = new RestTemplate();
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            if (config.getApiKey() != null && !config.getApiKey().isEmpty()) {
                headers.set("Authorization", "Bearer " + config.getApiKey());
            }
            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);
            ResponseEntity<String> response = restTemplate.postForEntity(config.getEndpoint(), entity, String.class);

            // 日志：响应体
            System.out.println("[Deepseek] Response: " + response.getBody());

            JSONObject json = JSON.parseObject(response.getBody());
            if (json.containsKey("choices")) {
                return json.getJSONArray("choices").getJSONObject(0).getJSONObject("message").getString("content");
            } else {
                return response.getBody();
            }
        } catch (HttpClientErrorException e) {
            String msg = errorCodeHandlerFactory
                .getHandler(MODEL_NAME)
                .getFriendlyMessage(e.getStatusCode().value(), e.getResponseBodyAsString());
            // throw new BusinessException(msg);
            return msg;
        }
    }

    @Override
    public Flux<String> chatStream(String prompt, Map<String, Object> params, AiModelConfigEntity config) {
        try {
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", config.getModelType() != null ? config.getModelType() : "deepseek-chat");
            List<Map<String, String>> messages = new ArrayList<>();
            if (params != null && params.containsKey("messages")) {
                messages = (List<Map<String, String>>) params.get("messages");
            } else {
                messages.add(Map.of("role", "user", "content", prompt));
            }
            requestBody.put("messages", messages);
            requestBody.put("stream", true);
            if (params != null) {
                requestBody.putAll(params);
            }

            WebClient webClient = WebClient.builder()
                .baseUrl(config.getEndpoint())
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
                .defaultHeader("Authorization", "Bearer " + config.getApiKey())
                .build();

            return webClient.post()
                .bodyValue(requestBody)
                .retrieve()
                .bodyToFlux(String.class)
                .filter(chunk -> chunk != null && !chunk.trim().isEmpty())
                .flatMap(chunk -> {
                    if (chunk.startsWith("data: ")) {
                        String data = chunk.substring(6);
                        if ("[DONE]".equals(data)) {
                            return Flux.empty();
                        }
                        try {
                            JSONObject json = JSON.parseObject(data);
                            if (json.containsKey("choices")) {
                                JSONObject choice = json.getJSONArray("choices").getJSONObject(0);
                                if (choice.containsKey("delta") && choice.getJSONObject("delta").containsKey("content")) {
                                    String content = choice.getJSONObject("delta").getString("content");
                                    return content != null ? Flux.just(content) : Flux.empty();
                                }
                            }
                        } catch (Exception e) {
                            // 忽略解析错误
                        }
                    }
                    return Flux.empty();
                })
                .onErrorResume(e -> {
                    String msg = "[AI服务连接失败] " + (e.getMessage() != null ? e.getMessage() : "未知错误");
                    return Flux.just(msg);
                });
        } catch (Exception e) {
            return Flux.just("[AI服务异常] " + (e.getMessage() != null ? e.getMessage() : "未知错误"));
        }
    }

    @Override
    public String getModelName() {
        return MODEL_NAME;
    }
} 